import { useEffect, useRef, useState } from 'preact/hooks'

export function useWebSocket(handlers) {
  const wsRef = useRef(null)
  const [isConnected, setIsConnected] = useState(false)
  const reconnectTimeoutRef = useRef(null)

  useEffect(() => {
    connectWebSocket()

    return () => {
      if (wsRef.current) {
        wsRef.current.close()
      }
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current)
      }
    }
  }, [])

  const connectWebSocket = () => {
    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
    const wsUrl = `${protocol}//${window.location.hostname}:8000/ws/graph-stream`

    const ws = new WebSocket(wsUrl)

    ws.onopen = () => {
      console.log('WebSocket connected')
      setIsConnected(true)
    }

    ws.onclose = () => {
      console.log('WebSocket disconnected')
      setIsConnected(false)

      // 尝试重连
      reconnectTimeoutRef.current = setTimeout(() => {
        console.log('Attempting to reconnect...')
        connectWebSocket()
      }, 3000)
    }

    ws.onerror = (error) => {
      console.error('WebSocket error:', error)
    }

    ws.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data)
        handleMessage(data)
      } catch (error) {
        console.error('Failed to parse WebSocket message:', error)
      }
    }

    wsRef.current = ws
  }

  const handleMessage = (data) => {
    const { type, payload, message } = data

    switch (type) {
      case 'processing_start':
        handlers.onProcessingStart?.()
        break

      case 'processing_complete':
        handlers.onProcessingComplete?.()
        break

      case 'analysis_complete':
        handlers.onAnalysisComplete?.(payload)
        break

      case 'node_added':
        if (payload && payload.node) {
          handlers.onNodeAdded?.(payload.node)
        }
        break

      case 'edge_added':
        if (payload && payload.edge) {
          handlers.onEdgeAdded?.(payload.edge)
        }
        break

      case 'topic_matched':
        if (payload) {
          handlers.onTopicMatched?.(payload)
        }
        break

      case 'recommendation_path':
        if (payload) {
          handlers.onRecommendationPath?.(payload)
        }
        break

      case 'counter_recommendations':
        if (payload && Array.isArray(payload)) {
          payload.forEach((rec) => {
            handlers.onRecommendationPath?.(rec)
          })
          handlers.onProcessingComplete?.()
        }
        break

      default:
        console.log('Unknown message type:', type)
    }
  }

  const sendMessage = (message) => {
    if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(message))
    } else {
      console.error('WebSocket is not connected')
    }
  }

  return {
    sendMessage,
    isConnected,
  }
}
